from torch import nn
import torch
import numpy as np
import sys
from sklearn.decomposition import PCA

# 循环平移
class LoopMove():
    def __init__(self, move):
        self.move = move

    def __call__(self, x):
        # x shape [bs len]
        x = np.roll(x, self.move, axis=1)
        return x

class RandomLoopMove():
    def __init__(self, min_move, max_move):
        self.min_move = min_move
        self.max_move = max_move

    def __call__(self, x):
        move = np.random.randint(self.min_move, self.max_move)
        loop_move = LoopMove(move)
        return loop_move(x)

# 整合, 将n个数整合成一个数，可以是平均值，最大值，最小值
# 然后进行邻近插值
class Integrate():
    def __init__(self, step, mode):
        self.step = step
        self.mode = mode

    def __call__(self, x):
        for i, data in enumerate(x):
            data = data.reshape(-1, self.step)
            ori_len = len(data)
            if self.mode == 'mean':
                data = np.mean(data, axis=1)
            elif self.mode == 'max':
                data = np.max(data, axis=1)
            elif self.mode == 'min':
                data = np.min(data, axis=1)
            for j, d in enumerate(data):
                x[i][2*j] = data[j]
                x[i][2*j+1] = data[j]
        return x

# 加入高斯噪声
class AddGaussianNoise():
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, x):
        x_tor = torch.from_numpy(x)
        noise = (torch.randn(x_tor[0].size()) * self.std + self.mean).numpy()
        for i, data in enumerate(x):
            x[i] = data + noise
        return x

# 水平翻转
class HorizontalFlip():
    def __init__(self):
        pass

    def __call__(self, x):
        for i, data in enumerate(x):
            x[i] = data[::-1]
        return x

# 标准化
class Normalize():
    def __init__(self):
        pass

    def __call__(self, x):
        for i, data in enumerate(x):
            data = (data - np.mean(data)) / np.std(data)
            x[i] = data
        return x
    
class PCATransform():
    def __init__(self, pca_dim):
        self.PCA = PCA(pca_dim)

    def __call__(self, x):
        x = self.PCA.fit_transform(x)
        return x

class Encoder():
    def __init__(self, encoder):
        self.encoder = encoder

    def __call__(self, x):
        x = self.encoder(x)
        return x

# class PCATransform2(nn.Module):
#     def __init__(self, pca_dim):
#         self.PCA = PCA(pca_dim)
# 
#     def __call__(self, x):
#         self.PCA.fit(x)
#         x = self.PCA.transform(x)
#         # 将x的维度保持和pca_dim一致
#         return x[:, -1 * pca_dim:]
